#include <strmif.h>
#include "cmemallocator.hpp"
#include "mediatypeutil.hpp"
#include <vfwmsgs.h>
#include <cassert>
#include <new>


HRESULT CMemAllocator::CreateInstance(IMemAllocator** pp)
{
    if (pp == 0)
        return E_POINTER;
        
    IMemAllocator*& p = *pp;
    
    p = new (std::nothrow) CMemAllocator();

    return p ? S_OK : E_OUTOFMEMORY;
}


CMemAllocator::CMemAllocator() :
    m_cRef(1),
    m_cActive(-1),  //means "properties not set"
    m_bCommitted(false),
    m_samples(0)
{
    m_hCond = CreateEvent(0, 0, 0, 0);
    assert(m_hCond);
    
    m_hMutex = CreateMutex(0, 0, 0);
    assert(m_hMutex);
}


CMemAllocator::~CMemAllocator()
{
    assert(m_cActive <= 0);
    assert(m_samples == 0);
    assert(!m_bCommitted);
    
    BOOL b = CloseHandle(m_hCond);
    assert(b);
    
    b = CloseHandle(m_hMutex);
    assert(b);
}


HRESULT CMemAllocator::QueryInterface(const IID& iid, void** ppv)
{
    if (ppv == 0)
        return E_POINTER;

    IUnknown*& pUnk = reinterpret_cast<IUnknown*&>(*ppv);
    
    if (iid == __uuidof(IUnknown))
        pUnk = this;
        
    else if (iid == __uuidof(IMemAllocator))
        pUnk = static_cast<IMemAllocator*>(this);
        
    else
    {
        pUnk = 0;
        return E_NOINTERFACE;
    }
            
    pUnk->AddRef();
    return S_OK;
}


ULONG CMemAllocator::AddRef()
{
    return InterlockedIncrement((LONG*)&m_cRef);
}


ULONG CMemAllocator::Release()
{
    if (LONG n = InterlockedDecrement((LONG*)&m_cRef))
        return n;
        
    delete this;
    return 0;
}


HRESULT CMemAllocator::SetProperties( 
    ALLOCATOR_PROPERTIES* pPreferred,
    ALLOCATOR_PROPERTIES* pActual)
{
    Lock lock;
        
    HRESULT hr = lock.Seize(this);

    if (FAILED(hr))
        return hr;
    
    if (m_bCommitted)
        return VFW_E_ALREADY_COMMITTED;    

    if (m_cActive > 0)
        return VFW_E_BUFFERS_OUTSTANDING;

    if (pPreferred == 0)  //weird
    {        
        m_props.cBuffers = 0;
        m_props.cbBuffer = 0;
        m_props.cbAlign = 1;
        m_props.cbPrefix = 0;        
    }
    else
    {
        m_props = *pPreferred;
        
        if (m_props.cBuffers < 0)
            m_props.cBuffers = 0;

        if (m_props.cbBuffer < 0)
            m_props.cbBuffer = 0;  

        if (m_props.cbAlign <= 0)
            m_props.cbAlign = 1;
            
        if (m_props.cbPrefix < 0)
            m_props.cbPrefix = 0;
    }
    
    if (pActual)
        *pActual = m_props;

    m_cActive = 0;  //means "properties have been set"
    
    return S_OK;
}


HRESULT CMemAllocator::GetProperties(ALLOCATOR_PROPERTIES* p)
{
    if (p == 0)
        return E_POINTER;
        
    Lock lock;
    
    HRESULT hr = lock.Seize(this);
        
    if (FAILED(hr))
        return hr;
    
    if (m_cActive < 0)
        return VFW_E_SIZENOTSET;

    *p = m_props;
    return S_OK;
}


HRESULT CMemAllocator::Commit()
{
    Lock lock;
    
    HRESULT hr = lock.Seize(this);
    
    if (FAILED(hr))
        return hr;
    
    if (m_bCommitted)
        return S_OK;

    if (m_cActive > 0)
        return VFW_E_BUFFERS_OUTSTANDING;
        
    if (m_cActive < 0)
        return VFW_E_SIZENOTSET;
        
    assert(m_samples == 0);
    
    for (long i = 0; i < m_props.cBuffers; ++i)
    {
        hr = allocate_sample();
        
        if (FAILED(hr))
            return hr;
    }
    
    m_bCommitted = true;

    return S_OK;
}


HRESULT CMemAllocator::Decommit()
{
    //NOTE: If there are any threads waiting on 
    //GetBuffer, then that op is supposed to 
    //finish, returning an error.
    //
    //Further calls to GetBuffer should fail, 
    //until Commit() is called.
    //
    //This op does *not* affect outstanding
    //buffers, so it's perfectly legal for
    //m_active to have a non-zero value.
    //
    //The purpose of this op is to prevent a
    //filter from getting any more samples.
    
    Lock lock;
    
    HRESULT hr = lock.Seize(this);
    
    if (FAILED(hr))
        return hr;
        
    m_bCommitted = false;
    
    while (m_samples)
    {
        Sample* const sample = m_samples;
        m_samples = sample->m_pNext;
        
        Sample::Destroy(sample);
    } 
    
    const BOOL b = SetEvent(m_hCond);
    assert(b);
    b;

    return S_OK;
}


HRESULT CMemAllocator::GetBuffer( 
    IMediaSample** pp,
    REFERENCE_TIME*,
    REFERENCE_TIME*,
    DWORD flags)
{
    if (pp == 0)
        return E_POINTER;
        
    Lock lock;
        
    if (flags & AM_GBF_NOWAIT)
    {
        const HRESULT hr = lock.Seize(this);
        
        if (FAILED(hr))
            return hr;
            
        if (!m_bCommitted)
            return VFW_E_NOT_COMMITTED;
            
        assert(m_cActive >= 0);
        
        if (m_cActive >= m_props.cBuffers)  //no samples available
            return VFW_E_TIMEOUT;
    }    
    else
    {
        for (;;)
        {
            HRESULT hr = lock.Seize(this);
            
            if (FAILED(hr))
                return hr;
            
            if (!m_bCommitted)
                return VFW_E_NOT_COMMITTED;
            
            assert(m_cActive >= 0);

            if (m_cActive < m_props.cBuffers)
                break;
                
            lock.Release();
            
            DWORD index;
            hr = CoWaitForMultipleHandles(
                    0, //wait all
                    INFINITE, 
                    1, 
                    &m_hCond, 
                    &index);
                                
            assert(hr == S_OK);
            assert(index == 0);
        }
    }
    
    IMediaSample*& p = *pp;    
    p = get_sample();
    
    AddRef();  //the contribution of this sample
    
    return S_OK;
}

    
HRESULT CMemAllocator::ReleaseBuffer(IMediaSample* p_)
{
    if (p_ == 0)
        return E_POINTER;
        
    Sample* p = static_cast<Sample*>(p_);
    //NOTE: Is there a way to make this type safe?
    //(One way is to use a custom interface that you
    //query for.)
        
    assert(p->m_cRef == 0);
    assert(p->m_pAllocator == this);
    
    Lock lock;

    const HRESULT hr = lock.Seize(this);
    
    if (FAILED(hr))
        return hr;
            
    assert(m_cActive > 0);
    --m_cActive;
    
    if (m_bCommitted)
    {
        p->Final();
        
        p->m_pNext = m_samples;
        m_samples = p;
    }
    else
    {
        Sample::Destroy(p);
    }
    
    const BOOL b = SetEvent(m_hCond);
    assert(b);
    b;
    
    //release lock now, in case this sample is holding 
    //the last reference to this allocator
    lock.Release();
    
    //now it's safe to release the allocator, after 
    //releasing the lock; if this causes the allocator
    //to be destroyed, then the lock won't attempt to
    //manipulate the allocator during its own destruction
    //(another thing to try would be to add a ref to the
    //allocator in the lock's ctor, and then release the
    //ref in its dtor.  This would prevent the destruction
    //of the allocator until the lock itself is destroyed).

    Release();  //the contribution of this sample
        
    return S_OK;
}


HRESULT CMemAllocator::allocate_sample()
{
    Sample* p;
    
    const HRESULT hr = Sample::Create(this, p);
    
    if (FAILED(hr))
        return hr;
        
    assert(p->m_cRef == 0);
    
    p->m_pNext = m_samples;
    m_samples = p;

    return S_OK;    
}


IMediaSample* CMemAllocator::get_sample()
{    
    assert(m_cActive >= 0);
    assert(m_cActive < m_props.cBuffers); 
    
    Sample* const p = m_samples;
    assert(p);
    assert(p->m_pAllocator == this);
    assert(p->m_cRef == 0);
    
    m_samples = p->m_pNext;
    p->m_pNext = 0;
    
    p->Init();
    
    ++m_cActive;

    return p;    
}



CMemAllocator::Lock::Lock() : m_hMutex(0)
{
}


CMemAllocator::Lock::~Lock()
{
    Release();
}


HRESULT CMemAllocator::Lock::Seize(CMemAllocator* p)
{
    assert(p);
    assert(m_hMutex == 0);

    //TODO:
    //We assume success here, but note that if the wait
    //were to fail, then you can't modify the lock
    //state -- so this assignment should probably come
    //later.
    m_hMutex = p->m_hMutex;
    
    //const DWORD dw = WaitForSingleObject(m_hMutex, INFINITE);
    //
    //if (dw == WAIT_FAILED)
    //{
    //    const DWORD e = GetLastError();
    //    return HRESULT_FROM_WIN32(e);
    //}
    //
    //assert(dw == WAIT_OBJECT_0);

    DWORD index;

    HRESULT hr = CoWaitForMultipleHandles(
                    0,  //wait flags
                    INFINITE, 
                    1, 
                    &m_hMutex, 
                    &index);
                        
    assert(hr == S_OK);
    assert(index == 0);
    hr;

    return S_OK;
}
    

void CMemAllocator::Lock::Release()
{
    if (m_hMutex)
    {
        const BOOL b = ReleaseMutex(m_hMutex);
        assert(b);
        b;
        
        m_hMutex = 0;
    }
}


HRESULT CMemAllocator::Sample::Create(
    CMemAllocator* pAllocator,
    Sample*& pSample)
{
    assert(pAllocator);

    const ALLOCATOR_PROPERTIES& props = pAllocator->m_props;
    assert(props.cbAlign >= 1);
    assert(props.cbPrefix >= 0);
    assert(props.cbBuffer >= 0);

    pSample = new (std::nothrow) Sample(pAllocator);
    
    if (pSample == 0)
        return E_OUTOFMEMORY;
        
    assert(pSample->m_cRef == 0);
    
    const long buflen = props.cbAlign - 1 + props.cbPrefix + props.cbBuffer;
    BYTE* const buf = new (std::nothrow) BYTE[buflen];
    
    if (buf == 0)
    {
        Destroy(pSample);
        pSample = 0;
        
        return E_OUTOFMEMORY;
    }
    
    pSample->m_buf = buf;
    pSample->m_buflen = buflen;       

    long off = props.cbPrefix;
    
    if (intptr_t n = intptr_t(buf) % props.cbAlign)
        off += props.cbAlign - n;

    pSample->m_off = off;
    
    //don't bother to Init here.  this is just to create a new 
    //sample to add to allocator's pool, during commit.  the sample will get
    //init'd later, during GetBuffer.
            
    return S_OK;
}


void CMemAllocator::Sample::Destroy(Sample* p)
{
    delete p;
}


CMemAllocator::Sample::Sample(CMemAllocator* p)
    : m_pAllocator(p),
      m_cRef(0),  //allocator will adjust 
      m_buf(0),
      m_buflen(0),
      m_off(0),
      m_pmt(0)
{
}


CMemAllocator::Sample::~Sample()
{
    Final();  //deallocate media type    
    delete[] m_buf;
}


void CMemAllocator::Sample::Init()
{
    assert(m_cRef == 0);
    assert(m_buf);
    assert(m_pmt == 0);
    
    m_sync_point = false;
    m_actual_data_length = 0;
    m_preroll = false;
    m_discontinuity = false;
    m_start_time = LLONG_MAX;
    m_stop_time = LLONG_MAX;
    m_media_start_time = LLONG_MAX;
    m_media_stop_time = LLONG_MAX;
    
    m_cRef = 1;  //ready to give to caller of GetBuffer
}


void CMemAllocator::Sample::Final()
{
    assert(m_cRef == 0);
    
    MediaTypeUtil::Free(m_pmt);
    m_pmt = 0;
}
    

HRESULT CMemAllocator::Sample::QueryInterface(const IID& iid, void** ppv)
{
    if (ppv == 0)
        return E_POINTER;
        
    IUnknown*& pUnk = reinterpret_cast<IUnknown*&>(*ppv);
    
    if (iid == __uuidof(IUnknown))
        pUnk = this;
        
    else if (iid == __uuidof(IMediaSample))
        pUnk = static_cast<IMediaSample*>(this);
        
    else
    {
        pUnk = 0;
        return E_NOINTERFACE;
    }
    
    pUnk->AddRef();
    return S_OK;
}


ULONG CMemAllocator::Sample::AddRef()
{
    return InterlockedIncrement((LONG*)&m_cRef);
}


ULONG CMemAllocator::Sample::Release()
{
    if (LONG n = InterlockedDecrement((LONG*)&m_cRef))
        return n;
    
    m_pAllocator->ReleaseBuffer(this);    
    return 0;
}        
    

HRESULT CMemAllocator::Sample::GetPointer(BYTE** pp)
{
    if (pp == 0)
        return E_POINTER;
        
    BYTE*& p = *pp;

    assert(m_buf);
    assert(m_buflen >= 0);
    assert(m_off >= 0);
    assert(m_off <= m_buflen);
    
    p = m_buf + m_off;

#ifdef _DEBUG
    const ALLOCATOR_PROPERTIES& props = m_pAllocator->m_props;
    assert(props.cbAlign >= 1);
    assert(props.cbPrefix >= 0);
    assert(intptr_t(p - props.cbPrefix) % props.cbAlign == 0);
#endif
    
    return S_OK;
}


long CMemAllocator::Sample::GetSize()
{
    assert(m_off <= m_buflen);

    const long size = m_buflen - m_off;

#ifdef _DEBUG
    const ALLOCATOR_PROPERTIES& props = m_pAllocator->m_props;
    assert(size >= props.cbBuffer);
#endif

    return size;
}


HRESULT CMemAllocator::Sample::GetTime( 
    REFERENCE_TIME* pstart,
    REFERENCE_TIME* pstop)
{
    if (m_start_time == LLONG_MAX)
        return VFW_E_SAMPLE_TIME_NOT_SET;
        
    if (pstart)
        *pstart = m_start_time;
        
    if (m_stop_time == LLONG_MAX)
        return VFW_S_NO_STOP_TIME;
        
    if (pstop)
        *pstop = m_stop_time;
        
    return S_OK;
}


HRESULT CMemAllocator::Sample::SetTime( 
    REFERENCE_TIME* pstart,
    REFERENCE_TIME* pstop)
{
    if (pstart)
        m_start_time = *pstart;
    else
        m_start_time = LLONG_MAX;
        
    if (pstop)
        m_stop_time = *pstop;
    else
        m_stop_time = LLONG_MAX;
        
    return S_OK;
}


HRESULT CMemAllocator::Sample::IsSyncPoint()
{
    return m_sync_point ? S_OK : S_FALSE;
}


HRESULT CMemAllocator::Sample::SetSyncPoint(BOOL b)
{
    m_sync_point = bool(b != 0);
    return S_OK;
}


HRESULT CMemAllocator::Sample::IsPreroll()
{
    return m_preroll ? S_OK : S_FALSE;
}


HRESULT CMemAllocator::Sample::SetPreroll(BOOL b)
{
    m_preroll = bool(b != 0);
    return S_OK;
}


long CMemAllocator::Sample::GetActualDataLength()
{
    return m_actual_data_length;
}


HRESULT CMemAllocator::Sample::SetActualDataLength(long len)
{
    m_actual_data_length = len;    
    return S_OK;
}


HRESULT CMemAllocator::Sample::GetMediaType( 
    AM_MEDIA_TYPE** pp)
{
    if (pp == 0)
        return E_POINTER;
        
    AM_MEDIA_TYPE*& p = *pp;
        
    if (m_pmt)
        return MediaTypeUtil::Create(*m_pmt, p);
        
    p = 0;
    return S_FALSE;  //means "no media type"
}


HRESULT CMemAllocator::Sample::SetMediaType( 
    AM_MEDIA_TYPE* pmt)
{
    if (pmt == 0)
    {
        MediaTypeUtil::Free(m_pmt);
        return S_OK;
    }
    
    return MediaTypeUtil::Create(*pmt, m_pmt);
}


HRESULT CMemAllocator::Sample::IsDiscontinuity()
{
    return m_discontinuity ? S_OK : S_FALSE;
}


HRESULT CMemAllocator::Sample::SetDiscontinuity(BOOL b)
{
    m_discontinuity = bool(b != 0);
    return S_OK;
}


HRESULT CMemAllocator::Sample::GetMediaTime( 
    LONGLONG* pstart,
    LONGLONG* pstop)
{
    if (m_media_start_time == LLONG_MAX)
        return VFW_E_MEDIA_TIME_NOT_SET;

    if (pstart)
        *pstart = m_media_start_time;
        
    if (pstop)
    {
        if (m_media_stop_time != LLONG_MAX)
            *pstop = m_media_stop_time;
        else
            *pstop = m_media_start_time + 1; //?
    }
        
    return S_OK;
}


HRESULT CMemAllocator::Sample::SetMediaTime( 
    LONGLONG* pstart,
    LONGLONG* pstop)
{
    if (pstart)
        m_media_start_time = *pstart;
    else
        m_media_start_time = LLONG_MAX;
        
    if (pstop)
        m_media_stop_time = *pstop;
    else
        m_media_stop_time = LLONG_MAX;

    return S_OK;
}
